import re
from typing import Dict, List, Tuple
from collections import Counter


class QuestionPreprocessor:
    """
    用户问题预处理类，用于纠错和改写问题，提升检索准确性
    """

    def __init__(self):
        """
        初始化问题预处理器
        """
        # 定义常见错别字映射
        self.typo_corrections = {
            # 常见错别字
            "流成图": "流程图",
            "流成程": "流程图",
            "过成图": "过程图",
            "过成": "过程",
            "检锁": "检索",
            "检所": "检索",
            "向俩": "向量",
            "向两": "向量",
            "余弦距里": "余弦距离",
            "余弦距": "余弦距离",
            "相俩": "向量",
            "相两": "向量",
            "相识度": "相似度",
            "相识性": "相似性",
            "模形": "模型",
            "模形化": "模型化",
            "向俩化": "向量化",
            "向两化": "向量化",
            "数库": "数据库",
            "据库": "数据库",
            "知库": "知识库",
            "识库": "知识库",
            "文当": "文档",
            "文旦": "文档",
            "文挡": "文档",
            "构键": "构建",
            "构健": "构建",
            "建健": "构建",
            "建键": "构建",
            "优话": "优化",
            "优划": "优化",
            "系通": "系统",
            "系疼": "系统",
            "原里": "原理",
            "原列": "原理",
            "工做": "工作",
            "实做": "实现",
            "实着": "实现",
            "技俩": "技术",
            "技两": "技术",
            "方发": "方法",
            "方乏": "方法",
            "规责": "规则",
            "规贼": "规则",
            "替欢": "替换",
            "替还": "替换",
            # 小说相关错别字
            # 小说相关词汇（无需纠正，但保留在字典中以避免误判）
            "修仙": "修仙",
            "修仙界": "修仙界",
            "穿越": "穿越",
            "玄幻": "玄幻",
            "都市": "都市",
            "科幻": "科幻",
            "悬疑": "悬疑",
            "仙侠": "仙侠",
            "历史": "历史",
            "奇幻": "奇幻",
            "武侠": "武侠",
            "言情": "言情",
            "恐怖": "恐怖",
            "现实": "现实",
            "斗气": "斗气",
            "武魂": "武魂",
            "废材": "废材",
            "逆袭": "逆袭",
            "系统": "系统",
            "爽文": "爽文",
            "重生": "重生",
            "转生": "转生",
            "科技": "科技",
            "创新": "创新",
            "惊悚": "惊悚",
            "异世界": "异世界",
            "机甲": "机甲",
            "星际": "星际",
            "战争": "战争",
            "神医": "神医",
            "兵王": "兵王",
            "秦朝": "秦朝",
            "三国": "三国",
            "明朝": "明朝",
            "侦探": "侦探",
            "魔法": "魔法",
            "龙族": "龙族",
            "精灵": "精灵",
            "剑仙": "剑仙",
            "剑客": "剑客",
            "江湖": "江湖",
            "总裁": "总裁",
            "甜宠": "甜宠",
            "校园": "校园",
            "青春": "青春",
            "娱乐圈": "娱乐圈",
            "鬼屋": "鬼屋",
            "诅咒": "诅咒",
            "奋斗": "奋斗",
            "打工人": "打工人",
            "创业": "创业",
            "境界": "境界",
            "功法": "功法",
            "丹药": "丹药",
            "法宝": "法宝",
            "门派": "门派",
            "职业": "职业",
            "神通": "神通",
            "阵法": "阵法",
            "霸总": "霸总",
            "情感线": "情感线",
            "场景": "场景",
            "世界观": "世界观",
            "异能": "异能",
            "怪物": "怪物",
            "朝代": "朝代",
            "服饰": "服饰",
            "饮食": "饮食",
            "建筑": "建筑",
            "礼仪": "礼仪",
            "案件": "案件",
            "侦查": "侦查",
            "反派": "反派",
            "宫斗": "宫斗",
            "权谋": "权谋",
            "种田": "种田",
            "经商": "经商",
            "人物": "人物",
            "情节": "情节",
            "梗": "梗",
            "描写": "描写",
            "对话": "对话",
            "叙述": "叙述",
        }
        
        # 定义问题改写规则
        self.rewrite_rules = [
            # 模糊问题到明确问题的映射
            (r"怎么用\s*RAG", "如何搭建基础RAG的知识库和检索流程"),
            (r"RAG\s*怎么", "RAG的工作原理和使用方法"),
            (r"什么是\s*RAG", "RAG的定义和核心组件"),
            (r"RAG\s*原理", "RAG的工作原理和关键技术"),
            (r"向量\s*搜索", "向量数据库检索原理和实现"),
            (r"向量\s*检索", "向量数据库检索原理和实现"),
            (r"知识库\s*构建", "如何构建RAG系统的知识库"),
            (r"文档\s*分割", "RAG系统中的文档分割策略"),
            (r"文档\s*切分", "RAG系统中的文档分割策略"),
            (r"文本\s*分割", "RAG系统中的文档分割策略"),
            (r"文本\s*切分", "RAG系统中的文档分割策略"),
            (r"向量\s*模型", "文本向量化模型的选择和应用"),
            (r"embedding\s*模型", "文本向量化模型的选择和应用"),
            (r"如何\s*优化\s*RAG", "RAG系统性能优化的方法和技巧"),
            (r"RAG\s*优化", "RAG系统性能优化的方法和技巧"),
            (r"向量\s*数据库", "向量数据库的选择和使用"),
            (r"Chroma\s*使用", "Chroma向量数据库的使用方法"),
            (r"知识\s*检索", "基于RAG的知识检索方法"),
            (r"语义\s*检索", "基于向量的语义检索技术"),
            # 小说相关问题改写
            (r"修仙\s*小说\s*创作", "穿越修仙类小说的创作要点和设定"),
            (r"玄幻\s*小说", "玄幻类小说的世界观构建和角色设定"),
            (r"都市\s*小说", "都市类小说的情节设计和人物塑造"),
            (r"科幻\s*小说", "科幻类小说的创意构思和技术设定"),
            (r"悬疑\s*小说", "悬疑类小说的剧情布局和推理设计"),
            (r"仙侠\s*小说", "仙侠类小说的修炼体系和门派设定"),
            (r"历史\s*小说", "历史类小说的背景还原和人物刻画"),
            (r"奇幻\s*小说", "奇幻类小说的种族设定和魔法体系"),
            (r"武侠\s*小说", "武侠类小说的武功招式和江湖背景"),
            (r"言情\s*小说", "言情类小说的情感线设计和人物关系"),
            (r"恐怖\s*小说", "恐怖类小说的氛围营造和惊悚元素"),
            (r"现实\s*小说", "现实类小说的题材选择和社会观察"),
            (r"斗气\s*大陆", "斗气大陆类小说的修炼等级和势力分布"),
            (r"系统\s*流", "系统流小说的系统功能和任务设计"),
            (r"废材\s*逆袭", "废材逆袭类小说的情节转折和成长路径"),
            (r"重生\s*类", "重生类小说的时间线设计和记忆优势"),
            (r"异世界\s*穿越", "异世界穿越类小说的世界观差异和适应过程"),
            (r"星际\s*战争", "星际战争类小说的科技设定和战略战术"),
            (r"校园\s*青春", "校园青春类小说的题材选择和情节设计"),
            (r"总裁\s*言情", "总裁言情类小说的人物设定和情节发展"),
            (r"鬼怪\s*恐怖", "鬼怪恐怖类小说的元素运用和氛围营造"),
            (r"职场\s*现实", "职场现实类小说的题材选择和情节设计"),
            (r"境界\s*体系", "小说中修炼境界体系的构建方法"),
            (r"功法\s*设定", "小说中功法技能的设定要点"),
            (r"丹药\s*设计", "小说中丹药道具的设计思路"),
            (r"法宝\s*构思", "小说中法宝装备的构思方法"),
            (r"门派\s*创建", "小说中门派势力的创建要点"),
            (r"职业\s*规划", "小说中角色职业的规划方法"),
            (r"神通\s*设计", "小说中神通技能的设计思路"),
            (r"阵法\s*设定", "小说中阵法体系的设定要点"),
            (r"情感线\s*设计", "小说中情感线的设计技巧"),
            (r"场景\s*描写", "小说中场景描写的技巧和方法"),
            (r"人物\s*塑造", "小说中人物塑造的方法和技巧"),
            (r"情节\s*构建", "小说中情节构建的技巧和方法"),
            (r"对话\s*写作", "小说中对话写作的技巧和要点"),
            (r"叙述\s*方式", "小说中叙述方式的选择和运用"),
        ]

    def correct_typos(self, question: str) -> str:
        """
        纠正问题中的错别字
        
        :param question: 原始问题
        :return: 纠错后的问题
        """
        corrected_question = question
        for typo, correct in self.typo_corrections.items():
            corrected_question = corrected_question.replace(typo, correct)
        return corrected_question

    def rewrite_question(self, question: str) -> str:
        """
        改写问题，将模糊问题明确化
        
        :param question: 原始问题
        :return: 改写后的问题
        """
        rewritten_question = question
        for pattern, replacement in self.rewrite_rules:
            if re.search(pattern, rewritten_question, re.IGNORECASE):
                rewritten_question = replacement
                break
        return rewritten_question

    @staticmethod
    def _question_similarity(q1: str, q2: str) -> float:
        """
        使用 Jaccard 相似度衡量两个问题的词重合度
        """
        if q1 == q2:
            return 1.0
        tokens1 = set(re.findall(r"[\u4e00-\u9fa5]+|[a-zA-Z0-9]+", q1.lower()))
        tokens2 = set(re.findall(r"[\u4e00-\u9fa5]+|[a-zA-Z0-9]+", q2.lower()))
        if not tokens1 or not tokens2:
            return 0.0
        intersection = len(tokens1 & tokens2)
        union = len(tokens1 | tokens2)
        return intersection / union if union else 0.0

    def preprocess(self, question: str) -> Dict[str, str]:
        """
        完整的问题预处理流程
        
        :param question: 原始问题
        :return: 预处理结果，包含原始问题、纠错后问题和改写后问题
        """
        # 1. 纠错
        corrected_question = self.correct_typos(question)
        
        # 2. 改写
        rewritten_question = self.rewrite_question(corrected_question)
        
        # 3. 返回处理结果
        return {
            "original": question,
            "corrected": corrected_question,
            "rewritten": rewritten_question
        }

    def add_typo_correction(self, typo: str, correction: str):
        """
        添加新的错别字纠正规则
        
        :param typo: 错别字
        :param correction: 正确词汇
        """
        self.typo_corrections[typo] = correction

    def add_rewrite_rule(self, pattern: str, replacement: str):
        """
        添加新的问题改写规则
        
        :param pattern: 匹配模式（正则表达式）
        :param replacement: 替换文本
        """
        self.rewrite_rules.append((pattern, replacement))

    def advanced_rewrite(self, question: str) -> str:
        """
        高级问题改写，基于关键词分析
        
        :param question: 原始问题
        :return: 改写后的问题
        """
        # 简单的关键词分割（替代jieba分词）
        words = re.findall(r'[\w]+', question)
        
        # 基于关键词组合进行改写
        if "RAG" in words and ("流程" in question or "步骤" in question):
            return "RAG系统构建的详细流程和关键步骤"
        
        if "向量" in words and ("相似度" in question or "距离" in question):
            return "向量相似度计算方法和距离度量"
            
        if "知识库" in words and ("更新" in question or "维护" in question):
            return "RAG知识库的维护和更新策略"
            
        if "模型" in words and ("选择" in question or "比较" in question):
            return "RAG系统中不同向量化模型的比较和选择标准"
        
        # 小说相关关键词改写
        if "修仙" in question and ("世界" in question or "体系" in question):
            return "穿越修仙类小说的世界观构建和修炼体系"
        
        if "玄幻" in question and ("大陆" in question or "世界" in question):
            return "玄幻类小说的大陆设定和种族分布"
        
        if "系统" in question and "任务" in question:
            return "系统流小说中的系统功能和任务设计"
        
        if "废材" in question and "逆袭" in question:
            return "废材逆袭类小说的角色成长路径和情节转折"
        
        if "重生" in question and ("前世" in question or "记忆" in question):
            return "重生类小说的时间线设计和记忆优势运用"
        
        if "异世界" in question and "穿越" in question:
            return "异世界穿越类小说的世界观差异和主角适应过程"
        
        if "星际" in question and "战争" in question:
            return "星际战争类小说的科技设定和战略战术"
        
        if "悬疑" in question and ("推理" in question or "侦探" in question):
            return "悬疑推理类小说的剧情布局和逻辑设计"
        
        if "言情" in question and ("情感" in question or "爱情" in question):
            return "言情类小说的情感线设计和人物关系处理"
        
        if "武侠" in question and ("江湖" in question or "武功" in question):
            return "武侠类小说的江湖背景和武功招式设定"
        
        if "历史" in question and ("朝代" in question or "背景" in question):
            return "历史类小说的朝代背景还原和人物刻画"
        
        if "恐怖" in question and ("氛围" in question or "惊悚" in question):
            return "恐怖类小说的氛围营造和惊悚元素运用"
        
        if "现实" in question and ("题材" in question or "社会" in question):
            return "现实类小说的题材选择和社会观察视角"
        
        if "总裁" in question and "言情" in question:
            return "总裁言情类小说的人物设定和情节发展"
        
        if "校园" in question and "青春" in question:
            return "校园青春类小说的题材选择和情节设计"
        
        if "境界" in question and "体系" in question:
            return "小说中修炼境界体系的构建方法和等级设定"
        
        if "功法" in question and "技能" in question:
            return "小说中功法技能的设定要点和表现形式"
        
        if "人物" in question and "塑造" in question:
            return "小说中人物塑造的方法和技巧"
        
        if "情节" in question and "构建" in question:
            return "小说中情节构建的技巧和方法"
        
        if "对话" in question and "写作" in question:
            return "小说中对话写作的技巧和要点"
        
        # 默认返回原问题
        return question

    def comprehensive_preprocess(self, question: str) -> Dict[str, str]:
        """
        综合问题预处理，包括基础和高级处理
        
        :param question: 原始问题
        :return: 预处理结果
        """
        # 基础预处理
        basic_result = self.preprocess(question)
        
        corrected = basic_result["corrected"]
        rewritten = basic_result["rewritten"]
        advanced_rewrite = self.advanced_rewrite(corrected)

        final_question = corrected

        if rewritten != corrected:
            similarity = self._question_similarity(corrected, rewritten)
            if similarity >= 0.8:
                final_question = rewritten

        if advanced_rewrite != final_question:
            similarity = self._question_similarity(final_question, advanced_rewrite)
            if similarity >= 0.8:
                final_question = advanced_rewrite
            else:
                advanced_rewrite = final_question
        
        return {
            "original": question,
            "corrected": corrected,
            "rewritten": rewritten,
            "advanced_rewrite": advanced_rewrite,
            "final": final_question
        }


# 使用示例
if __name__ == "__main__":
    # 创建问题预处理器实例
    preprocessor = QuestionPreprocessor()
    
    # 测试问题预处理
    test_questions = [
        "怎么用 RAG 流成图",
        "RAG 怎么实现向俩搜索",
        "什么是知识库构键",
        "如何优化 RAG 系统",
        "向量检锁的原理是什么",
        "向量相似度计算方法",
        "RAG知识库怎么更新"
    ]
    
    for question in test_questions:
        result = preprocessor.comprehensive_preprocess(question)
        print(f"原始问题: {result['original']}")
        print(f"纠错后: {result['corrected']}")
        print(f"基础改写: {result['rewritten']}")
        print(f"高级改写: {result['advanced_rewrite']}")
        print(f"最终问题: {result['final']}")
        print("-" * 50)